# image_mixer_node.py
from comfy.utils import ProgressBar
import torch

class ImageLatentMixer:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image1": ("IMAGE",),
                "image2": ("IMAGE",),
                "mix_factor": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}),
                "vae": ("VAE",),
            }
        }

    RETURN_TYPES = ("LATENT",)
    FUNCTION = "mix_images"
    CATEGORY = "latent/transforms"
    DESCRIPTION = "Blends two images in latent space using a mix factor."

    INPUT_TYPES_HINTS = {
        "mix_factor": { "widget": "slider" }
    }

    def mix_images(self, image1, image2, mix_factor, vae):
        # Encode both images into latent space
        latent1 = vae.encode(image1.movedim(-1, 1))  # Convert HxWxC -> CxHxW
        latent2 = vae.encode(image2.movedim(-1, 1))

        # Blend the latents
        mixed_latent = {
            "samples": latent1["samples"] * (1 - mix_factor) + latent2["samples"] * mix_factor
        }

        return (mixed_latent,)